{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "252e334b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "db0abc9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# MC MR MP\n",
    "NFR_TRANSITION = {\n",
    "    \"MST\": {\n",
    "(True, True, True): (0.9, 0.91, 0.9),\n",
    "(True, True, False): (0.88, 0.93, 0.85),\n",
    "(True, False, True): (0.92, 0.89, 0.92),\n",
    "(True, False, False): (0.9, 0.91, 0.87),\n",
    "(False, True, True): (0.85, 0.93, 0.88),\n",
    "(False, True, False): (0.83, 0.95, 0.83),\n",
    "(False, False, True): (0.87, 0.91, 0.9),\n",
    "(False, False, False): (0.85, 0.93, 0.85),\n",
    "    },\n",
    "    \"RT\": {\n",
    "(True, True, True): (0.86, 0.95, 0.82),\n",
    "(True, True, False): (0.84, 0.97, 0.75),\n",
    "(True, False, True): (0.88, 0.93, 0.84),\n",
    "(True, False, False): (0.86, 0.95, 0.77),\n",
    "(False, True, True): (0.73, 0.97, 0.8),\n",
    "(False, True, False): (0.71, 0.99, 0.73),\n",
    "(False, False, True): (0.75, 0.95, 0.82),\n",
    "(False, False, False): (0.73, 0.97, 0.75),      \n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "85571bdf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_states_dist(S):\n",
    "    # Get prob dist from S1 to S8\n",
    "    mc, mr, mp = S\n",
    "    dist = []\n",
    "    for i in [mc, 1-mc]:\n",
    "        for j in [mr, 1-mr]:\n",
    "            for k in [mp, 1-mp]:\n",
    "                dist.append(i * j * k)\n",
    "    return dist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "64c48125",
   "metadata": {},
   "outputs": [],
   "source": [
    "def P_NFR(S, A):\n",
    "    # S: Current Belief of states MC,MR,MP\n",
    "    # A: MST | RT\n",
    "    # Return the belief for the next state\n",
    "    dist = get_states_dist(S)\n",
    "    p_nfr = [0, 0, 0]\n",
    "    for i,p in enumerate(NFR_TRANSITION[A].values()):\n",
    "        for j in range(3):\n",
    "            p_nfr[j] += p[j] * dist[i]\n",
    "    return p_nfr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "f267dfe8",
   "metadata": {},
   "outputs": [],
   "source": [
    "OBS_RANGES = {\n",
    "    \"BC\": (0.7, 0.8),\n",
    "    \"ANL\": (0.08, 0.4),\n",
    "    \"TTW\": (0.7, 0.9),\n",
    "}\n",
    "\n",
    "OBS_TRANSITION = {\n",
    "    \"MST\": {\n",
    "        \"BC\": [0.8, 0.15, 0.05],\n",
    "        \"ANL\": [0.06, 0.16, 0.78],\n",
    "        \"TTW\": [0.83, 0.13, 0.04],\n",
    "    },\n",
    "    \"RT\":{\n",
    "        \"BC\": [0.78, 0.16, 0.06],\n",
    "        \"ANL\": [0.05, 0.15, 0.8],\n",
    "        \"TTW\": [0.8, 0.15, 0.05],\n",
    "    }\n",
    "}\n",
    "\n",
    "NS_OBS_TRANSITION = {\n",
    "    \"MST\": {\n",
    "        \"BC\": [0.72, 0.18, 0.1],\n",
    "        \"ANL\": [0.12, 0.2, 0.68],\n",
    "        \"TTW\": [0.67, 0.23, 0.1],\n",
    "    },\n",
    "    \"RT\":{\n",
    "        \"BC\": [0.68, 0.2, 0.12],\n",
    "        \"ANL\": [0.1, 0.18, 0.8],\n",
    "        \"TTW\": [0.63, 0.25, 0.12],\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "5235dd28",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_obs_ratio(obs, topo, obs_type):\n",
    "    if topo == \"MST\":\n",
    "        lb = 8 * 3\n",
    "        hb = 58 * 3\n",
    "    else:\n",
    "        lb = 50 * 3\n",
    "        hb = 83 * 3\n",
    "    if obs_type == \"BC\":\n",
    "        lb *= 20\n",
    "        hb *= 30\n",
    "    if obs_type == \"TTW\":\n",
    "        lb *= 10\n",
    "        hb *= 20\n",
    "    return (obs - lb) / (hb - lb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "fee0e7c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_obs_pos(obs, topo, obs_type):\n",
    "    pos = OBS_RANGES[obs_type]\n",
    "    obs_ratio = get_obs_ratio(obs, topo, obs_type)\n",
    "    if obs_ratio < pos[0]:\n",
    "        return 0\n",
    "    elif pos[0] <= obs_ratio < pos[1]:\n",
    "        return 1\n",
    "    else:\n",
    "        return 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "606c3379",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_obs_prob(obs, topo, obs_type):\n",
    "    pos = get_obs_pos(obs, topo, obs_type)\n",
    "    if pos == 0:\n",
    "        return OBS_RANGES[obs_type][0]\n",
    "    elif pos == 1:\n",
    "        return OBS_RANGES[obs_type][1] - OBS_RANGES[obs_type][0]\n",
    "    else:\n",
    "        return 1 - OBS_RANGES[obs_type][1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "9ceab3ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_obs_nfr_prob(obs, topo, obs_type):\n",
    "    pos = get_obs_pos(obs, topo, obs_type)\n",
    "    return OBS_TRANSITION[topo][obs_type][pos]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "fb7b7bfe",
   "metadata": {},
   "outputs": [],
   "source": [
    "def P_O_NFR(O, topo):\n",
    "    # O: bc, anl, ttw\n",
    "    bc, anl, ttw = O\n",
    "    #return [P_BC_MC(bc, topo), P_ANL_MR(anl, topo), P_TTW_MP(ttw, topo)]\n",
    "    return [\n",
    "        get_obs_nfr_prob(bc, topo, \"BC\"),\n",
    "        get_obs_nfr_prob(anl, topo, \"ANL\"),\n",
    "        get_obs_nfr_prob(ttw, topo, \"TTW\"),\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "e28be7f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def P_O_wrong(O, topo):\n",
    "    bc, anl, ttw = O\n",
    "    return [get_obs_prob(bc, topo, \"BC\"), get_obs_prob(anl, topo, \"ANL\"), get_obs_prob(ttw, topo, \"TTW\")]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "af604169",
   "metadata": {},
   "outputs": [],
   "source": [
    "obs_names = [\"BC\", \"ANL\", \"TTW\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "054f88ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_obs_dist(S, topo):\n",
    "    obs_trans = OBS_TRANSITION[topo]\n",
    "    ns_obs_trans = NS_OBS_TRANSITION[topo]\n",
    "    obs_dist = {\n",
    "        \"BC\": [0, 0, 0],\n",
    "        \"ANL\": [0, 0, 0],\n",
    "        \"TTW\": [0, 0, 0],\n",
    "    }\n",
    "    for i in range(3):\n",
    "        p_nfr = S[i]\n",
    "        for j in range(3):\n",
    "            obs_dist[obs_names[i]][j] = p_nfr * obs_trans[obs_names[i]][j] + (1 - p_nfr) * ns_obs_trans[obs_names[i]][j]\n",
    "    return obs_dist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "07b614a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def P_O(S, topo, O):\n",
    "    obs_dist = get_obs_dist(S, topo)\n",
    "    p_o = [0, 0, 0]\n",
    "    for i in range(3):\n",
    "        pos = get_obs_pos(O[i], topo, obs_names[i])\n",
    "        p_o[i] = obs_dist[obs_names[i]][pos]\n",
    "    return p_o"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "56d07ee5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def P_NFR_O(O, S, A):\n",
    "    p_n = P_NFR(S,A)\n",
    "    p_o_n = P_O_NFR(O,A)\n",
    "    p_o = P_O(S,A,O)\n",
    "    p = [0, 0, 0]\n",
    "    for i in range(3):\n",
    "        p[i] = p_n[i] * p_o_n[i] / p_o[i]\n",
    "    return p"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "4f93f3a5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.8263636363636364, 0.7483333333333334, 0.926867321867322]"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "P_NFR_O([6000, 50, 2000], [0.9, 0.2, 0.9], \"MST\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cd0a4b1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
